from ultralytics import YOLO
import cv2
import time 
nowTime=time.time()
lastTime=nowTime

yoloPath = "best.engine"
yolo11 = YOLO(yoloPath, task='classify')
# 4 th camera is realsense d435i RGB camera
cap = cv2.VideoCapture(4)
if not cap.isOpened():
    print("Error: Could not open camera.")
    exit()
predict_list = []
class_dict={0: '偏高', 1: '偏低', 2: '正常'}
class_dict_predict = {0: 0, 1:0, 2:0}
while True:
    ret, frame = cap.read()
    if not ret:
        print("Error: Could not read frame.")
        continue
    results = yolo11.predict(frame, imgsz=224, verbose=False)
    nowTime = time.time()
    for result in results:
        prob_top1 = result.probs.top1conf.item()
        # 筛选
        if prob_top1 < 0.9:
            continue
        lastTime=nowTime
        num_class_top1 = result.probs.top1
        #class_top1 = class_dict[num_class_top1]
        predict_list.append((num_class_top1, prob_top1))
    if nowTime - lastTime >2:
        #print("目前列表的长度为：",len(predict_list))
        predict_list.clear()
        #print("目前列表的长度为：",len(predict_list))
        lastTime=nowTime
    if len(predict_list) >= 10:
        # 统计
        for item in predict_list:
            class_dict_predict[item[0]] += 1
        max_class = max(class_dict_predict, key=class_dict_predict.get)
        # 清空计数 
        for name, value in class_dict_predict.items():
            class_dict_predict[name] = 0
        # 清空预测列表
        predict_list.clear()
        # 打印结果
        print("预测结果是",class_dict[max_class])
cap.release()
